# ----------------------------------------------------------------------------
# vectypes
# Copyright (c) 2009 Alex Holkner
# All rights reserved.
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions 
# are met:
#
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright 
#    notice, this list of conditions and the following disclaimer in
#    the documentation and/or other materials provided with the
#    distribution.
#  * Neither the name of vectypes nor the names of its
#    contributors may be used to endorse or promote products
#    derived from this software without specific prior written
#    permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
# ----------------------------------------------------------------------------

'''
Implements GLSL vector and matrix types, with some exceptions detailed below.
This implementation is based on the GLSL 1.20.8 specification, see
http://www.opengl.org/registry/doc/GLSLangSpec.Full.1.20.8.pdf

int, bool and float types from GLSL are substituted directly for
Python types of the same names.  As such, casting rules may differ from
the GLSL specification.

Types implemented:
  vec2,  vec3,  vec4
  ivec2, ivec3, ivec4
  bvec2, bvec3, bvec4

  mat2, mat2x2, mat2x3, mat2x4
  mat3, mat3x2, mat3x3, mat3x4, 
  mat4, mat4x2, mat4x3, mat4x4

Exceptions and additions to the specification:

  - Vectors can be constructed with no arguments; this is equivalent to
    initialising them from the (possibly truncated) vector vec4(0, 0, 0, 1).
  - Matrices can be constructed with no arguments, creating the identity matrix.
  - Vector components are accessible only in "xyzw" form.  The "rgba" and
    "stpq" forms are not available.
  - Only scalar attributes are accessible in the vector, for example "v.x".
    The vector and swizzzled forms, for example "v.xyz" and "v.xxx" are
    not supported.
  - Component indexing of vectors with indices less than zero is permitted and behaves
    as per the usual Python indexing (it addresses counting back from the end
    of the vector).  Additionally, slice indices can be used, and return
    a tuple of the components requested.

    For example:
      - vec2()[-1] is equivalent to vec2().y
      - v = vec2(); v[:] is equivalent to (v.x, v.y)

  - Division by zero raises ArithmeticException
  - True division and floor division are supported
  - Pre- and post-increment and decrement operators are not supported.
  - Positive unary operator (+) is supported in addition to unary negation (-)

'''

import math as _math

#: def vector_binop(op, vector_name, vector_component_type, vector_components):
        if not hasattr(other, '_vector_components'):
            # Component-wise scalar ___op___
            #: if vector_component_type == 'float':
            return ___vector_name___(
            #: else:
            return self._vector_upcast_scalar_rtype[type(other)](
            #: endif
                self.x ___op___ (other),
                self.y ___op___ (other),
            #:   if vector_components >= 3:
                self.z ___op___ (other),
            #:     if vector_components == 4:
                self.w ___op___ (other)
            #:     endif
            #:   endif
            )
        else:
            # Component-wise vector ___op___
            #: if vector_component_type == 'float':
            return ___vector_name___(
            #: else:
            return self._vector_upcast_vector_rtype[other._vector_base_type](
            #: endif
                self.x ___op___ (other.x),
                self.y ___op___ (other.y),
            #: if vector_components >= 3:
                self.z ___op___ (other.z),
            #:   if vector_components == 4:
                self.w ___op___ (other.w)
            #:   endif
            #: endif
            )
#: enddef

#: def vector_rev_scalar_binop(op, vector_name, vector_component_type, vector_components):
        # Component-wise scalar ___op___
        #: if vector_component_type == 'float':
        return ___vector_name___(
        #: else:
        return self._vector_upcast_scalar_rtype[type(other)](
        #: endif
            other ___op___ (self.x),
            other ___op___ (self.y),
        #:   if vector_components >= 3:
            other ___op___ (self.z),
        #:     if vector_components == 4:
            other ___op___ (self.w)
        #:     endif
        #:   endif
        )
#: enddef

#: def vector_unaryop(op, vector_name, vector_component_type, vector_components):
        # Component-wise scalar ___op___
        return ___vector_name___(
            ___op___ self.x,
            ___op___ self.y,
        #:   if vector_components >= 3:
            ___op___ self.z,
        #:     if vector_components == 4:
            ___op___ self.w
        #:     endif
        #:   endif
        )
#: enddef

#: def vector_eq(vector_name, vector_component_type, vector_components):
        if not hasattr(other, '_vector_components') or other._vector_components != self._vector_components:
            return False
        else:
            # Component-wise vector equality check
            return (self.x == other.x
                and self.y == other.y
            #: if vector_components >= 3:
                and self.z == other.z
            #:   if vector_components == 4:
                and self.w == other.w
            #:   endif
            #: endif
            )
#: enddef

#: def vector_ne(vector_name, vector_component_type, vector_components):
        if hasattr(other, '_vector_components') and other._vector_components == self._vector_components:
            # Component-wise vector equality check
            return (self.x != other.x
                 or self.y != other.y
            #: if vector_components >= 3:
                 or self.z != other.z
            #:   if vector_components == 4:
                 or self.w != other.w
            #:   endif
            #: endif
            )
        else:
            return True
#: enddef

#: # Helper methods used by vector * matrix
#: for n in (2, 3, 4):
def _dot___n___(a, b):
    # Implement dot product between two vectors of size ___n___
    return (
        a.x * b.x
      + a.y * b.y
    #: if n >= 3:
      + a.z * b.z
    #: endif
    #: if n == 4:
      + a.w * b.w
    #: endif
    )
#: endfor

#: def vector_outer_product(cols, rows):
#:  matrix_name = 'mat' + str(cols) + 'x' + str(rows)
            return ___matrix_name___(
              #: for col in range(cols):
                  #: for row in range(rows):
                    col[___col___] * row[___row___],
                  #: endfor
              #: endfor
            )
#: enddef

#: vector_types = []

#: def vector(vector_name, vector_component_type, vector_components):
#:     vector_base_type = vector_name[:-1]
def _unwrap___vector_name___args(args):
    if not args:
        # Identity constructor
        yield ___vector_component_type___(0)
        yield ___vector_component_type___(0)
        #: if vector_components >= 3:
        yield ___vector_component_type___(0)
        #:   if vector_components == 4:
        yield ___vector_component_type___(1)
        #:   endif
        #: endif

    elif len(args) == 1:
        arg = args[0]
        if hasattr(arg, '_vector_components'):
            # Copy constructor with optional truncation
            yield arg.x
            yield arg.y
            #: if vector_components >= 3:
            yield arg.z
            #:   if vector_components == 4:
            yield arg.w
            #:   endif
            #: endif
        else:
            # Unit scalar cast and replication
            #: for i in range(vector_components):
            yield ___vector_component_type___(arg)
            #: endfor

    else:
        # Initialize components sequentially from all arguments.
        for arg in args:
            argtype = type(arg)
            if argtype is ___vector_component_type___:
                # Fast path for common case -- scalar component
                yield arg 
            else:
                try:
                    components = arg._vector_components
                except AttributeError:
                    # Cast scalar component
                    yield ___vector_component_type___(arg)
                    continue

                # Component-wise cast components
                yield ___vector_component_type___(arg.x)
                yield ___vector_component_type___(arg.y)
                if components >= 3:
                    yield ___vector_component_type___(arg.z)
                    if components == 4:
                        yield ___vector_component_type___(arg.w)

class ___vector_name___(object):
    _vector_components = ___vector_components___
    #: vector_slots = 'xyzw'[:vector_components]
    __slots__ = tuple('___vector_slots___')

    _vector_base_type = '___vector_base_type___'

    def __init__(self, *args):
        #: if vector_components == 2:
        self.x, self.y = _unwrap___vector_name___args(args)
        #: elif vector_components == 3:
        self.x, self.y, self.z = _unwrap___vector_name___args(args)
        #: elif vector_components == 4:
        self.x, self.y, self.z, self.w = _unwrap___vector_name___args(args)
        #: endif

    def __repr__(self):
        return '___vector_name___%r' % (self[:],)

    def __getitem__(self, index):
        return (
            self.x,
            self.y,
        #: if vector_components >= 3:
            self.z,
        #:   if vector_components == 4:
            self.w,
        #:   endif
        #: endif
        )[index]

    def __setitem__(self, index, value):
        if index == 0:
            self.x = value
        elif index == 1:
            self.y = value
        #: if vector_components >= 3:
        elif index == 2:
            self.z = value
        #:   if vector_components == 4:
        elif index == 3:
            self.w = value
        #:   endif
        #: endif

    def __len__(self):
        return ___vector_components___

    def __add__(self, other):
        #: vector_binop('+', vector_name, vector_component_type, vector_components)

    def __sub__(self, other):
        #: vector_binop('-', vector_name, vector_component_type, vector_components)

    def __mul__(self, other):
        if hasattr(other, '_matrix_rows'):
            # row vector * matrix
            cols = other._cols
            assert other._matrix_rows == ___vector_components___, 'Vector and matrix must have compatible size'
            if other._matrix_cols == 2:
                return vec2(_dot___vector_components___(self, cols[0]),
                            _dot___vector_components___(self, cols[1]))
            elif other._matrix_cols == 3:
                return vec3(_dot___vector_components___(self, cols[0]),
                            _dot___vector_components___(self, cols[1]),
                            _dot___vector_components___(self, cols[2]))
            elif other._matrix_cols == 4:
                return vec4(_dot___vector_components___(self, cols[0]),
                            _dot___vector_components___(self, cols[1]),
                            _dot___vector_components___(self, cols[2]),
                            _dot___vector_components___(self, cols[3]))

        #: vector_binop('*', vector_name, vector_component_type, vector_components)

    def __div__(self, other):
        #: vector_binop('/', vector_name, vector_component_type, vector_components)

    def __truediv__(self, other):
        #: vector_binop('.__truediv__', vector_name, vector_component_type, vector_components)

    def __floordiv__(self, other):
        #: vector_binop('//', vector_name, vector_component_type, vector_components)

    def __radd__(self, other):
        #: vector_rev_scalar_binop('+', vector_name, vector_component_type, vector_components)

    def __rsub__(self, other):
        #: vector_rev_scalar_binop('-', vector_name, vector_component_type, vector_components)

    def __rmul__(self, other):
        #: vector_rev_scalar_binop('*', vector_name, vector_component_type, vector_components)

    def __rdiv__(self, other):
        #: vector_rev_scalar_binop('/', vector_name, vector_component_type, vector_components)

    def __rtruediv__(self, other):
        #: vector_rev_scalar_binop('.__truediv__', vector_name, vector_component_type, vector_components)

    def __rfloordiv__(self, other):
        #: vector_rev_scalar_binop('//', vector_name, vector_component_type, vector_components)

    def __neg__(self):
        #: vector_unaryop('-', vector_name, vector_component_type, vector_components) 

    def __pos__(self):
        #: vector_unaryop('+', vector_name, vector_component_type, vector_components) 

    def __eq__(self, other):
        #: vector_eq(vector_name, vector_component_type, vector_components)

    def __ne__(self, other):
        #: vector_ne(vector_name, vector_component_type, vector_components)

    def _outer_product(self, other):
        #: rows = vector_components
        col = self[:]
        row = other[:]
        cols = len(row)
        if cols == 2:
            #: vector_outer_product(2, rows)
        elif cols == 3:
            #: vector_outer_product(3, rows)
        elif cols == 4:
            #: vector_outer_product(4, rows)

#: vector_types.append((vector_name, vector_components, vector_base_type))
#: enddef

#: def vector_type_lookups():
#:     for vector_name, vector_components, vector_base_type in vector_types:
___vector_name___._vector_upcast_scalar_rtype = {
    float: vec___vector_components___,
#:          if vector_base_type == 'vec':
    int:   vec___vector_components___,
    bool:  vec___vector_components___,
#:          elif vector_base_type == 'ivec':
    int:   ivec___vector_components___,
    bool:  ivec___vector_components___,
#:          elif vector_base_type == 'bvec':
    int:   ivec___vector_components___,
    bool:  bvec___vector_components___,
#:          endif
}
___vector_name___._vector_upcast_vector_rtype = {
    'vec':   vec___vector_components___,
#:          if vector_base_type == 'vec':
    'ivec':  vec___vector_components___,
    'bvec':  vec___vector_components___,
#:          elif vector_base_type == 'ivec':
    'ivec':  ivec___vector_components___,
    'bvec':  ivec___vector_components___,
#:          elif vector_base_type == 'bvec':
    'ivec':  ivec___vector_components___,
    'bvec':  bvec___vector_components___,
#:          endif
}
#:     endfor
#: enddef

#: vector('vec2', 'float', 2)
#: vector('vec3', 'float', 3)
#: vector('vec4', 'float', 4)
#: vector('ivec2', 'int', 2)
#: vector('ivec3', 'int', 3)
#: vector('ivec4', 'int', 4)
#: vector('bvec2', 'bool', 2)
#: vector('bvec3', 'bool', 3)
#: vector('bvec4', 'bool', 4)

#: vector_type_lookups()

def _unwrap_matrix_args(args):
    # Initialize components sequentially from all arguments.
    for arg in args:
        argtype = type(arg)
        if argtype is float:
            # Fast path for common case -- scalar component
            yield arg 
        else:
            try:
                components = arg._vector_components
            except AttributeError:
                # Cast scalar component
                yield float(arg)
                continue

            # Component-wise cast components
            yield float(arg.x)
            yield float(arg.y)
            if components >= 3:
                yield float(arg.z)
                if components == 4:
                    yield float(arg.w)

#: def matrix_binop(op, matrix_name, matrix_cols, matrix_rows):
        if not hasattr(other, '_matrix_cols'):
            # Component-wise scalar ___op___
            return ___matrix_name___(
            #: for col in range(matrix_cols):
                self._cols[___col___] ___op___ (other),
            #: endfor
            )
        else:
            # Component-wise matrix ___op___
            return ___matrix_name___(
            #: for col in range(matrix_cols):
                self._cols[___col___] ___op___ (other._cols[___col___]),
            #: endfor
            )
#: enddef

#: def matrix_binop_scalar(op, matrix_name, matrix_cols, matrix_rows):
        # Component-wise scalar ___op___
        return ___matrix_name___(
        #: for col in range(matrix_cols):
            self._cols[___col___] ___op___ (other),
        #: endfor
        )
#: enddef

#: def matrix_rev_scalar_binop(op, matrix_name, matrix_cols, matrix_rows):
        # Component-wise rev scalar ___op___
        return ___matrix_name___(
        #: for col in range(matrix_cols):
            other ___op___ (self._cols[___col___]),
        #: endfor
        )
#: enddef

#: def matrix_unaryop(op, matrix_name, matrix_cols, matrix_rows):
        # Component-wise rev scalar ___op___
        return ___matrix_name___(
        #: for col in range(matrix_cols):
            ___op___ (self._cols[___col___]),
        #: endfor
        )
#: enddef

#: def matrix_mul(cols, rows, size):
#:   matrix_name = 'mat' + str(cols) + 'x' + str(rows)
                return ___matrix_name___(
                  #: for col in range(cols):
                  #:    for row in range(rows):
                      a[0][___row___] * b[___col___][0]
                  #:      for i in range(1, size):
                    + a[___i___][___row___] * b[___col___][___i___]
                  #:      endfor
                ,
                  #:    endfor
                  #: endfor
                )
#: enddef

#: def matrix(matrix_name, matrix_cols, matrix_rows):
#:   if matrix_rows == 2:
#:      matrix_column_type = 'vec2'
#:   elif matrix_rows == 3:
#:      matrix_column_type = 'vec3'
#:   elif matrix_rows == 4:
#:      matrix_column_type = 'vec4'
#:   endif

class ___matrix_name___(object):
    _matrix_cols = ___matrix_cols___
    _matrix_rows = ___matrix_rows___

    __slots__ = '_cols'

    def __init__(self, *args):
        if not args:
            self._init_diagonal(1.0)
        elif len(args) == 1:
            arg = args[0]
            if hasattr(arg, '_matrix_cols'):
                # Initialise from another matrix
                self._init_diagonal(1.0)
                for col in range(min(___matrix_cols___, arg._matrix_cols)):
                    arg_col = arg._cols[col]
                    self_col = self._cols[col]
                    for row in range(min(___matrix_rows___, arg._matrix_rows)):
                        self_col[row] = arg_col[row]
            else:
                self._init_diagonal(float(arg))
        else:
            components = list(_unwrap_matrix_args(args))
            #: i = 0
            #: j = i + matrix_rows
            self._cols = [
            #: for col in range(matrix_cols):
                ___matrix_column_type___(*components[___i___:___j___]),
            #:   i += matrix_rows
            #:   j += matrix_rows
            #: endfor
            ]

    def _init_diagonal(self, scalar):
        # Identity constructor
        self._cols = [
        #: for col in range(matrix_cols):
          ___matrix_column_type___(
        #:    for row in range(matrix_rows):
        #:       if row == col:
                    scalar,
        #:       else:
                    0.0,
        #:       endif
        #:    endfor
          ),
        #: endfor
        ]

    def __repr__(self):
        return '___matrix_name___%r' % (tuple(self._cols),)

    def pformat(self, indent=0):
        return '\n'.join([
        #: for row in range(matrix_rows):
            ' ' * indent + 
            '[ ' +
        #:   for col in range(matrix_cols):
                '%6.2f ' % self._cols[___col___][___row___] +
        #:   endfor
            ']',
        #: endfor
        ])

    def pprint(self, indent=0):
        print self.pformat(indent)

    def __getitem__(self, index):
        return self._cols[index]

    def __setitem__(self, index, value):
        self._cols[index] = value

    def __add__(self, other):
        #: matrix_binop('+', matrix_name, matrix_cols, matrix_rows)

    def __sub__(self, other):
        #: matrix_binop('-', matrix_name, matrix_cols, matrix_rows)

    def __mul__(self, other):
        if hasattr(other, '_vector_components'):
            # matrix * column vector
            assert other._vector_components == ___matrix_cols___
            return vec___matrix_rows___(
            #: for r in range(matrix_rows):
                  self._cols[0][___r___] * other.x
                + self._cols[1][___r___] * other.y
            #:  if matrix_cols >= 3:
                + self._cols[2][___r___] * other.z
            #:  endif
            #:  if matrix_cols == 4:
                + self._cols[3][___r___] * other.w
            #:  endif
                ,
            #: endfor
            )
        elif hasattr(other, '_matrix_rows'):
            # matrix * matrix
            assert other._matrix_rows == ___matrix_cols___
            a = self._cols
            b = other._cols
            if other._matrix_cols == 2:
                #: matrix_mul(2, matrix_rows, matrix_cols)
            elif other._matrix_cols == 3:
                #: matrix_mul(3, matrix_rows, matrix_cols)
            elif other._matrix_cols == 4:
                #: matrix_mul(4, matrix_rows, matrix_cols)

        #: matrix_binop_scalar('*', matrix_name, matrix_cols, matrix_rows)

    def _comp_mul(self, other):
        #: matrix_binop('*', matrix_name, matrix_cols, matrix_rows)

    def __div__(self, other):
        #: matrix_binop('/', matrix_name, matrix_cols, matrix_rows)

    def __truediv__(self, other):
        #: matrix_binop('.__truediv__', matrix_name, matrix_cols, matrix_rows)

    def __floordiv__(self, other):
        #: matrix_binop('//', matrix_name, matrix_cols, matrix_rows)

    def __radd__(self, other):
        #: matrix_rev_scalar_binop('+', matrix_name, matrix_cols, matrix_rows)

    def __rsub__(self, other):
        #: matrix_rev_scalar_binop('-', matrix_name, matrix_cols, matrix_rows)

    def __rmul__(self, other):
        #: matrix_rev_scalar_binop('*', matrix_name, matrix_cols, matrix_rows)

    def __rdiv__(self, other):
        #: matrix_rev_scalar_binop('/', matrix_name, matrix_cols, matrix_rows)

    def __rtruediv__(self, other):
        #: matrix_rev_scalar_binop('.__truediv__', matrix_name, matrix_cols, matrix_rows)

    def __rfloordiv__(self, other):
        #: matrix_rev_scalar_binop('//', matrix_name, matrix_cols, matrix_rows)

    def __neg__(self):
        #: matrix_unaryop('-', matrix_name, matrix_cols, matrix_rows) 
        pass

    def __pos__(self):
        #: matrix_unaryop('+', matrix_name, matrix_cols, matrix_rows) 
        pass

    def __eq__(self, other):
        return type(other) is type(self) and self._cols == other._cols

    def __ne__(self, other):
        return type(other) is not type(self) or self._cols != other._cols

    def _transpose(self):
        #: transpose_name = 'mat' + str(matrix_rows) + 'x' + str(matrix_cols)
        cols = self._cols
        return ___transpose_name___(
            #: for row in range(matrix_rows):
                #: for col in range(matrix_cols):
                   cols[___col___][___row___],
                #: endfor
            #: endfor
        )

#: enddef


#: matrix('mat2x2', 2, 2)
mat2 = mat2x2
#: matrix('mat2x3', 2, 3)
#: matrix('mat3x3', 3, 3)
#: matrix('mat3x2', 3, 2)
mat3 = mat3x3
#: matrix('mat2x4', 2, 4)
#: matrix('mat3x4', 3, 4)
#: matrix('mat4x4', 4, 4)
#: matrix('mat4x3', 4, 3)
#: matrix('mat4x2', 4, 2)
mat4 = mat4x4

# Vector functions

def length(v):
    n = v._vector_components
    if n == 2:
        return _math.sqrt(v.x * v.x + v.y * v.y)
    elif n == 3:
        return _math.sqrt(v.x * v.x + v.y * v.y + v.z * v.z)
    elif n == 4:
        return _math.sqrt(v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w)

def distance(a, b):
    return length(a - b)

def dot(a, b):
    '''Dot product of vectors a and b'''
    assert a._vector_components == b._vector_components, 'Vectors must have equal size'
    n = a._vector_components
    if n == 2:
        return (a.x * b.x
              + a.y * b.y)
    elif n == 3:
        return (a.x * b.x
              + a.y * b.y
              + a.z * b.z)
    else:
        return (a.x * b.x
              + a.y * b.y
              + a.z * b.z
              + a.w * b.w)

def cross(a, b):
    return vec3(
        a.y * b.z - b.y * a.z,
        a.z * b.x - b.z * a.x,
        a.x * b.y - b.x * a.y)

def normalize(v):
    return v / length(v)

def faceforward(N, I, Nref):
    if dot(Nref, I) < 0:
        return N
    else:
        return -N

def reflect(I, N):
    return I - 2 * dot(N, I) * N

def refract(I, N, eta):
    d = dot(N, I)
    k = 1.0 - eta * eta * (1.0 - d * d)
    if k < 0.0:
        return I.__class__()
    else:
        return eta * I - (eta * d + _math.sqrt(k)) * N

# Matrix functions
def matrixCompMult(a, b):
    return a._comp_mul(b)

def outerProduct(c, r):
    return c._outer_product(r)

def transpose(m):
    return m._transpose()
